import pandas as pd
import numpy as np
import json



def write_json(result,file_name):
    with open(file_name,'w',encoding='utf-8')as f:
        json.dump(result,f,ensure_ascii=False)


def build4leves(df,json_file):
    """ 最早的一款，分四层 """
    # levels = ['领域', '技术定位', '一级能力', '二级能力']
    levels = ['Unnamed: 1', '技术名称', '一级能力', '二级能力']
    # levels = ['field', 'Technical positioning', 'Level 1 capability', 'Level 2 capability']
    colors = {'空白技术':'blue','共研技术':'pink','通用技术':'#99e64d','nan':'grey'}

    result = []
    # 领域
    level0 = list(set(df[levels[0]]))
    for field in level0:
        # df[df['领域'] == '控制']
        df0 = df[df[levels[0]] == field]
        item0 = {'name':field, 'value': df0.shape[0], 'children':[]}

        # 技术定位 -> 技术名称
        tp_list = list(set(df0[levels[1]]))
        for tp in tp_list:
            tp_nan = 'blank' if tp is np.nan else tp
            # df0[df0['技术定位'] == '共研技术']
            df1 = df0[df0[levels[1]] == tp]
            item1 = {'name':tp_nan, 'value': df1.shape[0],}
            # item1['color'] = colors.get(tp,'grey')
            # item1['itemStyle'] = {"color":colors.get(tp,'grey')}
            item1['children'] = []

            # 一级能力
            c1_list = list(set(df0[levels[2]]))
            for c1 in c1_list:
                # df0[df0['一级能力'] == 'IT搭建能力']
                df2 = df0[df0[levels[2]] == c1]
                item2 = {'name':c1, 'value': df2.shape[0],}
                # 颜色
                if df2[df2['技术定位']=='空白技术'].shape[0] > 0:
                    color = colors['空白技术']
                elif df2[df2['技术定位']=='共研技术'].shape[0] > 0:
                    color = colors['共研技术']
                else:
                    color = colors['通用技术']
                item2['itemStyle'] = {"color":color}

                item2['children'] = []

                # 二级能力
                c2_list = list(set(df2[levels[3]]))
                for c2 in c2_list:
                    # df3 = df2[df2['二级能力']=='网络部署能力']
                    df3 = df2[df2[levels[3]]==c2]
                    item3 = {'name':c2, 'value': 1, 'depth':4,}
                    # item3['color'] = colors.get(df3['技术定位'].values[0],'grey')
                    item3['itemStyle'] = {"color":colors.get(df3['技术定位'].values[0],'grey')}


                    item2['children'].append(item3)
                item1['children'].append(item2)
            item0['children'].append(item1)
        result.append(item0)

    write_json(result,json_file)


def cal_value(df,rate=1):
    value = 0
    tp_pool = df['技术定位'].tolist()
    for tp in tp_pool:
        if tp == '空白技术':
            value += 5
        elif tp == '共研技术':
            value += 1
        else:
            value += 1
    return value*rate


def build_tp(df):
    """ 以技术定位为颜色区分 """
    # levels = ['领域', '技术定位', '一级能力', '二级能力']
    levels = ['Unnamed: 1', '一级能力', '二级能力']
    # levels = ['field', 'Technical positioning', 'Level 1 capability', 'Level 2 capability']

    colors = {'空白技术':'blue','共研技术':'pink','通用技术':'#99e64d','nan':'grey'}

    result = []
    # 领域
    level0 = list(set(df[levels[0]]))
    rate = 1
    for field in level0:
        if field in ['氢硼']:
            rate = 1.5
        else:
            rate = 1
        # df[df['领域'] == '控制']
        df0 = df[df[levels[0]] == field]
        item0 = {'name':field, 'value':cal_value(df0,rate),'num': df0.shape[0], 'children':[]}

        # 一级能力
        c1_list = list(set(df0[levels[1]]))
        for c1 in c1_list:
            # df0[df0['一级能力'] == 'IT搭建能力']
            df2 = df0[df0[levels[1]] == c1]
            item2 = {'name':c1, 'value':cal_value(df2,rate),'num': df2.shape[0],}
            # 颜色
            if df2[df2['技术定位']=='空白技术'].shape[0] > 0:
                color = colors['空白技术']
            elif df2[df2['技术定位']=='共研技术'].shape[0] > 0:
                color = colors['共研技术']
            else:
                color = colors['通用技术']
            item2['itemStyle'] = {"color":color}

            item2['children'] = []

            # 二级能力
            c2_list = list(set(df2[levels[2]]))
            for c2 in c2_list:
                # df3 = df2[df2['二级能力']=='网络部署能力']
                df3 = df2[df2[levels[2]]==c2]
                item3 = {'name':c2, 'value':cal_value(df3,rate), 'num': 1,}
                # item3['color'] = colors.get(df3['技术定位'].values[0],'grey')
                item3['itemStyle'] = {"color":colors.get(df3['技术定位'].values[0],'grey')}
                # '组织现状', '需求时间','建设思路'
                item3['组织现状'] = df3['组织现状'].values[0]
                item3['需求时间'] = df3['需求时间'].values[0]
                item3['建设思路'] = df3['建设思路'].values[0]


                item2['children'].append(item3)
            item0['children'].append(item2)
        result.append(item0)

    return result


def build_oc(df):
    """ 以组织现状为颜色区分 """
    # levels = ['领域', '技术定位', '一级能力', '二级能力']
    levels = ['Unnamed: 1', '一级能力', '二级能力']
    # levels = ['field', 'Technical positioning', 'Level 1 capability', 'Level 2 capability']
    # 组织现状
    colors = {'无相关能力':'blue','有相关能力，但能力不足':'pink','有相关能力，但资源不足':'yellow','有相关能力':'#99e64d','nan':'grey'}

    result = []
    # 领域
    level0 = list(set(df[levels[0]]))
    rate = 1
    for field in level0:
        # if field in ['物理','加热','氢硼']: rate = 2
        if field in ['氢硼']:
            rate = 1.5
        else:
            rate = 1
        # df[df['领域'] == '控制']
        df0 = df[df[levels[0]] == field]
        item0 = {'name':field, 'value':cal_value(df0,rate),'num': df0.shape[0], 'children':[]}

        # 一级能力
        c1_list = list(set(df0[levels[1]]))
        for c1 in c1_list:
            # df0[df0['一级能力'] == 'IT搭建能力']
            df2 = df0[df0[levels[1]] == c1]
            item2 = {'name':c1, 'value':cal_value(df2,rate),'num': df2.shape[0],}
            # 颜色
            if df2[df2['组织现状']=='无相关能力'].shape[0] > 0:
                color = colors['无相关能力']
            elif df2[df2['组织现状']=='有相关能力，但能力不足'].shape[0] > 0:
                color = colors['有相关能力，但能力不足']
            elif df2[df2['组织现状']=='有相关能力，但资源不足'].shape[0] > 0:
                color = colors['有相关能力，但资源不足']
            else:
                color = colors['有相关能力']
            item2['itemStyle'] = {"color":color}

            item2['children'] = []

            # 二级能力
            c2_list = list(set(df2[levels[2]]))
            for c2 in c2_list:
                # df3 = df2[df2['二级能力']=='网络部署能力']
                df3 = df2[df2[levels[2]]==c2]
                item3 = {'name':c2, 'value':cal_value(df3,rate), 'num': 1,}
                # item3['color'] = colors.get(df3['技术定位'].values[0],'grey')
                item3['itemStyle'] = {"color":colors.get(df3['组织现状'].values[0],'grey')}
                # '组织现状', '需求时间','建设思路'
                item3['需求时间'] = df3['需求时间'].values[0]
                item3['建设思路'] = df3['建设思路'].values[0]


                item2['children'].append(item3)
            item0['children'].append(item2)
        result.append(item0)

    return result


colors = {'空无':'blue','空有':'#4169E1','共无':'#FFa6b1','共有':'pink','通无':'#88e64d','通有':'#aaff5e','nan':'grey'}
def c1_color(df):
    if df[df['技术定位']=='空白技术'].shape[0] > 0:
        color = colors['空无']
    elif df[df['技术定位']=='共研技术'].shape[0] > 0:
        color = colors['共无']
    else:
        color = colors['通无']
    return color

def c2_color(df):
    tp = df['技术定位'].values[0]
    oc = df['组织现状'].values[0]
    color = 'grey'
    if tp=='空白技术' and (oc=='无相关能力' or oc=='有相关能力，但资源不足'):
        color = colors['空无']
    elif tp=='空白技术' and (oc=='有相关能力' or oc=='有相关能力，但能力不足'):
        color = colors['空有']
    elif tp=='共研技术' and (oc=='无相关能力' or oc=='有相关能力，但资源不足'):
        color = colors['共无']
    elif tp=='共研技术' and (oc=='有相关能力' or oc=='有相关能力，但能力不足'):
        color = colors['共有']
    elif tp=='通用技术' and (oc=='无相关能力' or oc=='有相关能力，但资源不足'):
        color = colors['通无']
    elif tp=='通用技术' and (oc=='有相关能力' or oc=='有相关能力，但能力不足'):
        color = colors['通有']

    return color


def build_co(df):
    """ 以技术定位和组织现状为颜色区分 """
    # levels = ['领域', '技术定位', '一级能力', '二级能力']
    levels = ['Unnamed: 1', '一级能力', '二级能力']
    # levels = ['field', 'Technical positioning', 'Level 1 capability', 'Level 2 capability']

    result = []
    # 领域
    level0 = list(set(df[levels[0]]))
    rate = 1
    for field in level0:
        rate = 1.5 if field in ['氢硼'] else 1

        # df[df['领域'] == '控制']
        df0 = df[df[levels[0]] == field]
        item0 = {'name':field, 'value':cal_value(df0,rate),'num': df0.shape[0], 'children':[]}

        # 一级能力
        c1_list = list(set(df0[levels[1]]))
        for c1 in c1_list:
            # df0[df0['一级能力'] == 'IT搭建能力']
            df2 = df0[df0[levels[1]] == c1]
            item2 = {'name':c1, 'value':cal_value(df2,rate),'num': df2.shape[0],}
            item2['itemStyle'] = {"color":c1_color(df2)}
            item2['children'] = []

            # 二级能力
            c2_list = list(set(df2[levels[2]]))
            for c2 in c2_list:
                # df3 = df2[df2['二级能力']=='网络部署能力']
                df3 = df2[df2[levels[2]]==c2]
                item3 = {'name':c2, 'value':cal_value(df3,rate), 'num': 1,}
                item3['itemStyle'] = {"color":c2_color(df3)}
                # 技术定位 '组织现状', '需求时间','建设思路'
                item3['技术定位'] = df3['技术定位'].values[0]
                item3['组织现状'] = df3['组织现状'].values[0]
                item3['需求时间'] = df3['需求时间'].values[0]
                item3['建设思路'] = df3['建设思路'].values[0]


                item2['children'].append(item3)
            item0['children'].append(item2)
        result.append(item0)

    return result



if __name__ == "__main__":
    # excel_name = 'data/聚变能力地图初稿202304.xlsx'
    # excel_name = 'data/聚变能力地图初稿202304v2.xlsx'
    excel_name = 'data/聚变能力地图初稿202304v4.xlsx'
    df = pd.read_excel(excel_name,header=2,sheet_name='能力地图基础数据')
    # avail_columns = ['Unnamed: 0','领域', '技术名称', '技术定位', '定位分类', '一级能力','二级能力','行业现状', '组织现状', '需求时间','建设思路',]
    avail_columns = ['Unnamed: 0','Unnamed: 1', '技术名称', '技术定位', '定位分类', '一级能力','二级能力','行业现状', '组织现状', '需求时间','建设思路',]
    useless_columns = [col for col in df.columns if col not in avail_columns]
    df.drop(columns=useless_columns, inplace=True)

    # result = build_tp(df)
    # write_json(result,'data/abilityMap_tp.json')

    # result = build_oc(df)
    # write_json(result,'data/abilityMap_oc.json')

    result = build_co(df)
    write_json(result,'data/abilityMap.json')

    time_list = list(set(df['需求时间']))
    current_list = [t for t in time_list if '2023' in str(t)]
    df23 = df[df['需求时间'].isin(current_list)]

    result = build_co(df23)
    write_json(result,'data/abilityMap_23.json')


print('finish')